"""
数据库迁移脚本：为 workflow_node_executions 表添加异步执行所需字段

执行方式：
python migrations/001_add_async_fields_to_node_execution.py
"""
import sqlite3
import os
from pathlib import Path


def migrate():
    """执行数据库迁移"""
    # 获取数据库路径
    backend_dir = Path(__file__).parent.parent
    db_path = backend_dir / "data" / "app.db"

    if not db_path.exists():
        print(f"❌ 数据库文件不存在: {db_path}")
        return False

    print(f"📊 连接数据库: {db_path}")
    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()

    try:
        # 检查表是否存在
        cursor.execute("""
            SELECT name FROM sqlite_master
            WHERE type='table' AND name='workflow_node_executions'
        """)
        if not cursor.fetchone():
            print("❌ 表 workflow_node_executions 不存在")
            return False

        # 检查字段是否已存在
        cursor.execute("PRAGMA table_info(workflow_node_executions)")
        columns = [row[1] for row in cursor.fetchall()]

        # 需要添加的字段
        new_columns = {
            'external_task_id': 'VARCHAR(255)',
            'submitted_at': 'TIMESTAMP',
            'poll_count': 'INTEGER DEFAULT 0',
            'next_poll_at': 'TIMESTAMP',
            'poll_interval': 'INTEGER DEFAULT 300'
        }

        # 添加缺失的字段
        added_columns = []
        for column_name, column_type in new_columns.items():
            if column_name not in columns:
                print(f"➕ 添加字段: {column_name} {column_type}")
                cursor.execute(f"""
                    ALTER TABLE workflow_node_executions
                    ADD COLUMN {column_name} {column_type}
                """)
                added_columns.append(column_name)
            else:
                print(f"⏭️  字段已存在，跳过: {column_name}")

        # 创建索引以优化查询
        print("📇 创建索引...")

        # 检查索引是否存在
        cursor.execute("""
            SELECT name FROM sqlite_master
            WHERE type='index' AND name='idx_node_exec_status_poll'
        """)
        if not cursor.fetchone():
            cursor.execute("""
                CREATE INDEX idx_node_exec_status_poll
                ON workflow_node_executions(status, next_poll_at)
            """)
            print("✅ 创建索引: idx_node_exec_status_poll")
        else:
            print("⏭️  索引已存在: idx_node_exec_status_poll")

        cursor.execute("""
            SELECT name FROM sqlite_master
            WHERE type='index' AND name='idx_workflow_instance_status'
        """)
        if not cursor.fetchone():
            cursor.execute("""
                CREATE INDEX idx_workflow_instance_status
                ON workflow_instances(status)
            """)
            print("✅ 创建索引: idx_workflow_instance_status")
        else:
            print("⏭️  索引已存在: idx_workflow_instance_status")

        # 提交更改
        conn.commit()

        if added_columns:
            print(f"\n✅ 迁移成功！添加了 {len(added_columns)} 个字段: {', '.join(added_columns)}")
        else:
            print("\n✅ 迁移完成，所有字段已存在，无需更改")

        # 验证字段
        cursor.execute("PRAGMA table_info(workflow_node_executions)")
        print("\n📋 当前表结构:")
        for row in cursor.fetchall():
            print(f"  - {row[1]} ({row[2]})")

        return True

    except Exception as e:
        conn.rollback()
        print(f"\n❌ 迁移失败: {str(e)}")
        return False

    finally:
        conn.close()


def rollback():
    """回滚迁移（SQLite不支持DROP COLUMN，需要重建表）"""
    backend_dir = Path(__file__).parent.parent
    db_path = backend_dir / "data" / "app.db"

    if not db_path.exists():
        print(f"❌ 数据库文件不存在: {db_path}")
        return False

    print(f"⚠️  警告：SQLite 不支持 DROP COLUMN，回滚需要重建表")
    print(f"建议使用数据库备份恢复")
    return False


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "rollback":
        rollback()
    else:
        migrate()
